#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多媒体文件处理器
处理SVGA、Lottie、WebP、音视频等文件的指纹修改
"""

import os
import json
import hashlib
import random
import shutil
import tempfile
from pathlib import Path
from typing import Dict, Optional
from zipfile import ZipFile, ZIP_DEFLATED
from datetime import datetime
from ..builtin_config_module.builtin_config import BuiltInConfig
from ..utils_module.logger import ObfuscationLogger

# 导入必要的库
try:
    from PIL import Image
    PILLOW_AVAILABLE = True
except ImportError:
    PILLOW_AVAILABLE = False

try:
    import mutagen
    from mutagen.mp3 import MP3
    from mutagen.mp4 import MP4
    from mutagen.id3 import ID3, TIT2, TPE1, COMM
    MUTAGEN_AVAILABLE = True
except ImportError:
    MUTAGEN_AVAILABLE = False

try:
    from fontTools.ttLib import TTFont
    FONTTOOLS_AVAILABLE = True
except ImportError:
    FONTTOOLS_AVAILABLE = False

class MultimediaHandler:
    """多媒体文件处理器"""
    
    def __init__(self, config: BuiltInConfig, logger: ObfuscationLogger):
        self.config = config
        self.logger = logger
        
        # 支持的文件格式
        self.supported_formats = {
            'image': ['.webp', '.apng', '.ico', '.bmp', '.tiff'],
            'animation': ['.svga', '.lottie', '.json'],  # lottie是json格式
            'audio': ['.mp3', '.m4a', '.aac', '.wav', '.ogg', '.flac'],
            'video': ['.mp4', '.mov', '.avi', '.webm', '.mkv'],
            'font': ['.ttf', '.otf', '.woff', '.woff2'],
        }
    
    def modify_file_fingerprint(self, file_path: str) -> bool:
        """修改文件指纹"""
        try:
            ext = Path(file_path).suffix.lower()
            
            # 根据文件类型选择处理方法
            if ext == '.svga':
                return self._modify_svga_file(file_path)
            elif ext == '.lottie' or (ext == '.json' and self._is_lottie_file(file_path)):
                return self._modify_lottie_file(file_path)
            elif ext == '.webp':
                return self._modify_webp_file(file_path)
            elif ext in self.supported_formats['audio']:
                return self._modify_audio_file(file_path)
            elif ext in self.supported_formats['video']:
                return self._modify_video_file(file_path)
            elif ext in self.supported_formats['font']:
                return self._modify_font_file(file_path)
            else:
                return False
                
        except Exception as e:
            self.logger.log_error(f"修改文件失败 {file_path}: {e}")
            return False
    
    def _modify_svga_file(self, file_path: str) -> bool:
        """修改SVGA文件 - SVGA是protobuf格式"""
        tmp_path = None
        try:
            tmp_fd, tmp_path = tempfile.mkstemp(suffix='.svga')
            os.close(tmp_fd)

            metadata = {
                'modified': True,
                'fingerprint': hashlib.sha256(os.urandom(16)).hexdigest(),
                'timestamp': datetime.utcnow().isoformat() + 'Z'
            }

            with ZipFile(file_path, 'r') as src_zip, ZipFile(tmp_path, 'w', ZIP_DEFLATED) as dst_zip:
                for info in src_zip.infolist():
                    data = src_zip.read(info.filename)
                    dst_zip.writestr(info, data)

                dst_zip.writestr(
                    'obfuscator_meta.json',
                    json.dumps(metadata, ensure_ascii=False, indent=2)
                )

            shutil.move(tmp_path, file_path)
            self.logger.log_operation("SVGA文件修改完成", file_path)
            return True

        except Exception as exc:
            self.logger.log_error(f"修改SVGA文件失败: {exc}")
            return self._append_random_padding(file_path, b'SVGAOBF')
        finally:
            if tmp_path and os.path.exists(tmp_path):
                try:
                    os.remove(tmp_path)
                except OSError:
                    pass
    
    def _is_lottie_file(self, file_path: str) -> bool:
        """检查JSON文件是否是Lottie动画"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            # Lottie文件通常包含这些字段
            return 'v' in data and ('layers' in data or 'assets' in data)
        except:
            return False
    
    def _modify_lottie_file(self, file_path: str) -> bool:
        """修改Lottie动画文件 - JSON格式"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 添加自定义元数据
            data['_metadata'] = {
                'modified': True,
                'fingerprint': hashlib.md5(os.urandom(16)).hexdigest()[:8]
            }
            
            # 如果有nm(name)字段，可以微调
            if 'nm' in data:
                data['nm'] = data['nm'] + ' '  # 添加空格不影响显示
            
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, separators=(',', ':'))
            
            self.logger.log_operation("Lottie文件修改完成", file_path)
            return True
            
        except Exception as e:
            self.logger.log_error(f"修改Lottie文件失败: {e}")
            return False
    
    def _modify_webp_file(self, file_path: str) -> bool:
        """修改WebP文件"""
        if not PILLOW_AVAILABLE:
            self.logger.log_warning("Pillow库未安装，无法处理WebP文件")
            return False
        
        try:
            # 使用Pillow处理WebP
            img = Image.open(file_path)
            
            # 检查是否为动画
            is_animated = getattr(img, 'is_animated', False)
            
            if is_animated:
                # 对于动画WebP，保留所有帧
                frames = []
                durations = []
                
                # 获取帧数
                n_frames = getattr(img, 'n_frames', 1)
                
                for i in range(n_frames):
                    img.seek(i)
                    frames.append(img.copy())
                    # 获取帧持续时间
                    duration = img.info.get('duration', 100)
                    durations.append(duration)
                
                # 修改第一帧的信息
                frames[0].info['comment'] = f"Modified_{random.randint(10000, 99999)}"
                
                # 保存动画，保持所有帧
                frames[0].save(
                    file_path,
                    format='WEBP',
                    save_all=True,
                    append_images=frames[1:],
                    duration=durations,
                    loop=img.info.get('loop', 0),
                    background=img.info.get('background', (255, 255, 255, 255)),
                    quality=random.randint(95, 100),
                    method=6
                )
            else:
                # 静态WebP
                # WebP支持EXIF和XMP元数据
                if hasattr(img, 'info'):
                    img.info['comment'] = f"Modified_{random.randint(10000, 99999)}"
                
                # 微调质量重新保存
                quality = random.randint(95, 100)
                img.save(file_path, format='WEBP', quality=quality, method=6)
            
            self.logger.log_operation("WebP文件修改完成", file_path)
            return True
            
        except Exception as e:
            self.logger.log_error(f"修改WebP文件失败: {e}")
            return False
    
    def _modify_audio_file(self, file_path: str) -> bool:
        """修改音频文件元数据"""
        if not MUTAGEN_AVAILABLE:
            # 如果没有mutagen，尝试写入ID3v1标签
            wrote_tag = self._write_id3v1_tag(file_path)
            if wrote_tag:
                self.logger.log_operation("音频标签写入", file_path)
                return True
            # 退回到简单方法
            return self._append_random_padding(file_path, b'ID3FALLBACK')
        
        try:
            # 使用mutagen修改音频元数据
            audio = mutagen.File(file_path)
            
            if audio is None:
                return self._append_random_padding(file_path, b'AUDNULL')
            
            # 添加或修改注释
            random_comment = f"Modified_{random.randint(10000, 99999)}"
            
            if isinstance(audio, MP3):
                # MP3使用ID3标签
                if audio.tags is None:
                    audio.add_tags()
                audio.tags.add(COMM(encoding=3, text=random_comment))
            elif isinstance(audio, MP4):
                # MP4使用不同的标签系统
                audio.tags['©cmt'] = random_comment
            else:
                # 通用处理
                if hasattr(audio, 'tags') and audio.tags:
                    audio.tags['comment'] = random_comment
            
            audio.save()
            self.logger.log_operation("音频文件修改完成", file_path)
            return True
            
        except Exception as e:
            self.logger.log_error(f"修改音频文件失败: {e}")
            if self._write_id3v1_tag(file_path):
                return True
            return self._append_random_padding(file_path, b'AUDERR')

    def _modify_video_file(self, file_path: str) -> bool:
        """修改视频文件元数据"""
        if not MUTAGEN_AVAILABLE:
            return self._append_random_padding(file_path, b'VIDFALL')
        
        try:
            # MP4视频可以用mutagen处理
            if file_path.lower().endswith('.mp4'):
                video = MP4(file_path)
                video.tags['©cmt'] = f"Modified_{random.randint(10000, 99999)}"
                video.save()
                self.logger.log_operation("视频文件修改完成", file_path)
                return True
            else:
                # 其他格式追加随机块，保持容错
                return self._append_random_padding(file_path, b'VIDMOD')
                
        except Exception as e:
            self.logger.log_error(f"修改视频文件失败: {e}")
            return self._append_random_padding(file_path, b'VIDERR')
    
    def _modify_font_file(self, file_path: str) -> bool:
        """修改字体文件"""
        try:
            if FONTTOOLS_AVAILABLE and file_path.lower().endswith(('.ttf', '.otf', '.woff', '.woff2')):
                if self._modify_font_with_fonttools(file_path):
                    self.logger.log_operation("字体元数据更新", file_path)
                    return True
            # 回退到在末尾添加安全填充
            return self._append_random_padding(file_path, b'FONTMOD')
        except Exception as e:
            self.logger.log_error(f"修改字体文件失败: {e}")
            return False

    def _append_random_padding(self, file_path: str, marker: bytes) -> bool:
        """在文件末尾安全地追加随机数据"""
        try:
            padding = marker + os.urandom(random.randint(8, 32))
            with open(file_path, 'ab') as handle:
                handle.write(padding)
            return True
        except Exception as exc:
            self.logger.log_warning(f"追加随机填充失败 {file_path}: {exc}")
            return False

    def _write_id3v1_tag(self, file_path: str) -> bool:
        """为MP3追加ID3v1标签以改变指纹"""
        try:
            # ID3v1结构: TAG + title + artist + album + year + comment + genre
            tag = bytearray(128)
            tag[0:3] = b'TAG'
            label = f"Mod{random.randint(1000, 9999)}".encode('latin-1', errors='ignore')
            tag[3:3+len(label[:30])] = label[:30]
            artist = b'Obfuscator'
            tag[33:33+len(artist[:30])] = artist[:30]
            comment = os.urandom(28)
            tag[97:97+len(comment)] = comment
            tag[127] = 0  # genre undefined

            with open(file_path, 'rb+') as handle:
                handle.seek(-128, os.SEEK_END)
                existing = handle.read(128)
                if existing.startswith(b'TAG'):
                    handle.seek(-128, os.SEEK_END)
                else:
                    handle.seek(0, os.SEEK_END)
                handle.write(bytes(tag))
            return True
        except Exception as exc:
            self.logger.log_warning(f"写入ID3v1标签失败 {file_path}: {exc}")
            return False

    def _modify_font_with_fonttools(self, file_path: str) -> bool:
        """使用fontTools更新字体名称表"""
        try:
            font = TTFont(file_path)
            name_table = font['name']
            new_suffix = f"-obf{random.randint(100, 999)}"
            for name_id in (1, 4):  # Font Family, Full Name
                entries = name_table.getName(name_id, 3, 1, 0x409)
                if entries:
                    for record in entries if isinstance(entries, list) else [entries]:
                        value = record.toUnicode()
                        record.string = (value + new_suffix).encode('utf-16-be')
            tmp_path = file_path + '.tmp'
            font.save(tmp_path)
            font.close()
            os.replace(tmp_path, file_path)
            return True
        except Exception as exc:
            self.logger.log_warning(f"fontTools 处理字体失败 {file_path}: {exc}")
            return False
    
    def batch_process(self, directory: str) -> Dict[str, int]:
        """批量处理目录下的多媒体文件"""
        stats = {
            'total': 0,
            'success': 0,
            'failed': 0
        }
        
        # 收集所有支持的扩展名
        all_extensions = []
        for exts in self.supported_formats.values():
            all_extensions.extend(exts)
        
        # 遍历目录
        for root, dirs, files in os.walk(directory):
            for file in files:
                ext = Path(file).suffix.lower()
                if ext in all_extensions:
                    file_path = os.path.join(root, file)
                    stats['total'] += 1
                    
                    if self.modify_file_fingerprint(file_path):
                        stats['success'] += 1
                    else:
                        stats['failed'] += 1
        
        return stats
